import numpy as np
import ctypes
import cv2
import os
from pykinect2 import PyKinectV2
from pykinect2.PyKinectV2 import *
from pykinect2 import PyKinectRuntime
from PIL import Image

kinect = PyKinectRuntime.PyKinectRuntime(PyKinectV2.FrameSourceTypes_Depth | PyKinectV2.FrameSourceTypes_Color)

file_dir = "D:\\RGBD_Face_datasets\\Lock3DFace-share\\color"
img_list = []
for root, dirs, files in os.walk(file_dir):
    if files:
        for file in files:
            img_path = root
            img_list.append(root + "\\" + file)
            img_path = img_path.replace("color", "calibrate")
            # print(img_path)
            # print(os.path.exists(img_path))
            if not os.path.exists(img_path):
                os.makedirs(img_path)

depth_path = "D:\\RGBD_Face_datasets\\Lock3DFace-share\\depth\\001_Kinect_FE_1DEPTH\\01.png"

for color_path in img_list:
    color_img = cv2.imread(color_path)
    depth_img = Image.open(depth_path)
    align_color_img = np.array(depth_img)

    depth2color_points_type = _ColorSpacePoint * np.int(512 * 424)
    depth2color_points = ctypes.cast(depth2color_points_type(), ctypes.POINTER(_ColorSpacePoint))
    kinect.mapper.MapDepthFrameToColorSpace(ctypes.c_uint(512 * 424), kinect.depth_frame_data,
                                            kinect.depth_frame_data_capacity, depth2color_points)
    colorXYs = np.copy(np.ctypeslib.as_array(depth2color_points,
                                             shape=(kinect.depth_frame_desc.Height * kinect.depth_frame_desc.Width,)))
    colorXYs = colorXYs.view(np.float32).reshape(colorXYs.shape + (-1,))
    colorXYs += 0.5
    colorXYs = colorXYs.reshape(kinect.depth_frame_desc.Height, kinect.depth_frame_desc.Width, 2).astype(np.int)
    colorXs = np.clip(colorXYs[:, :, 0], 0, kinect.color_frame_desc.Width - 1)
    colorYs = np.clip(colorXYs[:, :, 1], 0, kinect.color_frame_desc.Height - 1)
    new_img = color_img[colorYs, colorXs, :]

    # print(align_color_img.shape)
    # print(new_img.shape)

    # cv2.imshow('img', cv2.flip(align_color_img, 1))
    # cv2.imshow('img', new_img)
    color_path = color_path.replace("color", "calibrate")
    cv2.imwrite(color_path, new_img)
    cv2.waitKey(0)